{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JGKVzg-YIoMo"
      },
      "source": [
        "##### Copyright 2019 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3xtKPbtuIrwK"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6uTJAWnKIz6G"
      },
      "source": [
        "# High Performance Monte Carlo Simulation of Ising Model on TPU Clusters\n",
        "\n",
        "This notebook is a companion webpage for the paper: *High Performance Monte Carlo Simulation of Ising Model on TPU Clusters (Yang et al., 2019)*. See the [README.md](https://github.com/google-research/google-research/blob/master/simulation_research/ising_model/README.md) for details on how to simulate Ising model on Cloud TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cOu1dN5gEQmN"
      },
      "outputs": [],
      "source": [
        "\"\"\"Ising Model MCMC Simulation on TPU.\n",
        "\n",
        "This is the implementation of Algorithm 2: UpdateOptim in the paper.\n",
        "\"\"\"\n",
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import itertools\n",
        "import os\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow.contrib.tpu.python.ops import tpu_ops\n",
        "from tensorflow.contrib.cluster_resolver import TPUClusterResolver\n",
        "from tensorflow.python.ops import inplace_ops\n",
        "\n",
        "# Constants\n",
        "_NANOS_PER_SECOND = 1e9\n",
        "\n",
        "# Whether bfloat16 is used in the simulation.\n",
        "USE_BFLOAT16 = True\n",
        "# Lattice sub dimension, where each sub lattice is a square grid of spin values.\n",
        "LATTICE_SUB_DIM = 256\n",
        "# Lattice block dimensions of size 2, where each block is a grid of sub lattices\n",
        "# with shape lattice_block_dims.\n",
        "# Note: The TPU chips accessible from https://colab.research.google.com are\n",
        "# TPUv2 (8 cores). They have less memory than TPUv3 (the chips used in the\n",
        "# paper) so can not fit [224, 112], the size reported in the paper.\n",
        "# `README.md` has instructions on setting up Google Cloud project to use\n",
        "# TPUv3 and running larger experiments.\n",
        "LATTICE_BLOCK_DIMS = [96, 48]\n",
        "# Lattice super block dimension, where each super block is a square grid of\n",
        "# lattice blocks.\n",
        "LATTICE_SUPER_BLOCK_DIM = 2\n",
        "# Cores topology of size 2, each core is one replica of the computation.\n",
        "CORES_TOPOLOGY = [2, 4]\n",
        "# Burn in steps in MCMC.\n",
        "NUMBER_OF_BURN_IN_UPDATE = 1\n",
        "# Simulation steps in MCMC.\n",
        "NUMBER_OF_WHOLE_LATTICE_UPDATE = 200\n",
        "# Critical temperature.\n",
        "CRITICAL_TEMPERATURE = 2.26918531421\n",
        "# The temperature normalized to the critical temerature. Default to 1.0.\n",
        "NORMALIZED_TEMPERATURE = 1.0\n",
        "\n",
        "# Using Convolution for nearest neighbor computation.\n",
        "USE_CONV = True\n",
        "# Convolution kernel channel sizes.\n",
        "KERNEL_CHANNEL_SIZE = 16\n",
        "KERNEL_CHANNEL_SIZE_2 = 128\n",
        "CONV_LATTICE_DIMS = [(LATTICE_SUB_DIM // 2) * dim for dim in LATTICE_BLOCK_DIMS]\n",
        "\n",
        "\n",
        "# END PARAMETERS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PK9s-J8AHgQK"
      },
      "source": [
        "###TensorFlow Version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 336,
          "status": "ok",
          "timestamp": 1573969016180,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "qU0iSJglHeMd",
        "outputId": "3ae6dfc4-4d61-4b0b-f4fa-b4bbc6ac20bc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "1.15.0\n"
          ]
        }
      ],
      "source": [
        "print(tf.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DdoYRHLcHwNQ"
      },
      "outputs": [],
      "source": [
        "def get_inverse_temperature():\n",
        "  \"\"\"Returns the inverse temperature.\"\"\"\n",
        "  return 1.0 / (NORMALIZED_TEMPERATURE * CRITICAL_TEMPERATURE)\n",
        "\n",
        "def grid_coordinates(computation_shape):\n",
        "  \"\"\"Returns a numpy array containing all grid coordinates.\n",
        "\n",
        "  Args:\n",
        "    computation_shape: A sequence of integers giving the shape of the grid.\n",
        "\n",
        "  Returns:\n",
        "    A numpy array with shape\n",
        "    (np.prod(computation_shape), len(computation_shape)) and type np.int32.\n",
        "  \"\"\"\n",
        "  rank = len(computation_shape)\n",
        "  assert rank \u003e 0\n",
        "  coords = np.meshgrid(\n",
        "      *[np.arange(x, dtype=np.int32) for x in computation_shape], indexing=\"ij\")\n",
        "  return np.stack(coords, axis=-1).reshape(-1, rank)\n",
        "\n",
        "def tpu_device_assignment(computation_shape, tpu_topology):\n",
        "  \"\"\"Builds a DeviceAssignment that maps grid coordinates to TPU cores.\"\"\"\n",
        "  # This may be too restrictive, but it makes mapping onto the TPU topology\n",
        "  # simple. We can use a more complicated algorithm if needed.\n",
        "  if computation_shape.ndim != 1:\n",
        "    raise ValueError(\n",
        "        \"computation_shape ({}) must be a vector\".format(computation_shape))\n",
        "  if np.prod(computation_shape) \u003e np.prod(tpu_topology.mesh_shape):\n",
        "    raise ValueError(\n",
        "        \"computation_shape ({}) does not fit in TPU mesh shape ({})\".format(\n",
        "            computation_shape, tpu_topology.mesh_shape))\n",
        "\n",
        "  core_assignment = grid_coordinates(\n",
        "      tpu_topology.mesh_shape)[0:np.prod(computation_shape), :]\n",
        "  core_assignment = core_assignment.reshape((-1, 1, tpu_topology.mesh_rank))\n",
        "  # External logical compute grid shape.\n",
        "  compute_core_assignment = grid_coordinates(computation_shape)\n",
        "  return (tf.contrib.tpu.DeviceAssignment(tpu_topology, core_assignment),\n",
        "          compute_core_assignment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZbAibb596lna"
      },
      "outputs": [],
      "source": [
        "def get_dtype():\n",
        "  return tf.bfloat16 if USE_BFLOAT16 else tf.float32\n",
        "\n",
        "\n",
        "def is_single_core(cores_topology):\n",
        "  return cores_topology == [1, 1]\n",
        "\n",
        "\n",
        "def create_iterator(shape):\n",
        "  \"\"\"Create an iterator with a given shape.\"\"\"\n",
        "  dims = [range(dim) for dim in shape]\n",
        "  return itertools.product(*dims)\n",
        "\n",
        "\n",
        "def create_list(shape):\n",
        "  \"\"\"Create a list with a given shape and default value None.\"\"\"\n",
        "  if shape:\n",
        "    return [create_list(shape[1:]) for _ in range(shape[0])]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Yxo66RithMHK"
      },
      "outputs": [],
      "source": [
        "class NearestNeighborCalculatorOptimConv(object):\n",
        "  \"\"\"Calculate the sum of nearest neighbor spin values.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    # Constant matrix to compute the sum of nearest neighbor spins. These are\n",
        "    # formulated as the kernel, or filters, for tf.nn.conv1d and when applied,\n",
        "    # it computes the sum of nearest neighbors.\n",
        "    self._kernel_nn = tf.constant(np.expand_dims(np.stack(\n",
        "        [np.eye(KERNEL_CHANNEL_SIZE, k=-15, dtype=np.float16),\n",
        "         (np.eye(KERNEL_CHANNEL_SIZE, k=0, dtype=np.float16) +\n",
        "          np.eye(KERNEL_CHANNEL_SIZE, k=1, dtype=np.float16)),\n",
        "         np.zeros([KERNEL_CHANNEL_SIZE, KERNEL_CHANNEL_SIZE],\n",
        "                  dtype=np.float16)]), axis=0), tf.bfloat16)\n",
        "    self._kernel_nn_t = tf.constant(np.expand_dims(np.stack(\n",
        "        [np.zeros([KERNEL_CHANNEL_SIZE, KERNEL_CHANNEL_SIZE], dtype=np.float16),\n",
        "         (np.eye(KERNEL_CHANNEL_SIZE, k=-1, dtype=np.float16) +\n",
        "          np.eye(KERNEL_CHANNEL_SIZE, k=0, dtype=np.float16)),\n",
        "         np.eye(KERNEL_CHANNEL_SIZE, k=15, dtype=np.float16)]), axis=0),\n",
        "         tf.bfloat16)\n",
        "    self._kernel_nn_2 = tf.constant(np.expand_dims(np.stack(\n",
        "        [np.eye(KERNEL_CHANNEL_SIZE_2, k=-15, dtype=np.float16),\n",
        "         (np.eye(KERNEL_CHANNEL_SIZE_2, k=0, dtype=np.float16) +\n",
        "          np.eye(KERNEL_CHANNEL_SIZE_2, k=1, dtype=np.float16)),\n",
        "         np.zeros([KERNEL_CHANNEL_SIZE_2, KERNEL_CHANNEL_SIZE_2],\n",
        "                  dtype=np.float16)]), axis=0), tf.bfloat16)\n",
        "    self._kernel_nn_t_2 = tf.constant(np.expand_dims(np.stack(\n",
        "        [np.zeros([KERNEL_CHANNEL_SIZE_2, KERNEL_CHANNEL_SIZE_2],\n",
        "                  dtype=np.float16),\n",
        "         (np.eye(KERNEL_CHANNEL_SIZE_2, k=-1, dtype=np.float16) +\n",
        "          np.eye(KERNEL_CHANNEL_SIZE_2, k=0, dtype=np.float16)),\n",
        "         np.eye(KERNEL_CHANNEL_SIZE_2, k=15, dtype=np.float16)]), axis=0),\n",
        "         tf.bfloat16)\n",
        "    # If the lattice is distributed in multiple replicas, then define that\n",
        "    # permutation pairs that permute the boundaries of sub-lattices across\n",
        "    # replicas in all 4 directions. Those boundary spins are used to compute the\n",
        "    # nearest neighor sums of boundary spins on sub-lattices. They are ignored\n",
        "    # if the lattice is updated on single core.\n",
        "    tpu_x, tpu_y = CORES_TOPOLOGY\n",
        "    core_ids = np.arange(tpu_x * tpu_y).reshape((tpu_x, tpu_y))\n",
        "    core_ids_n = np.roll(core_ids, 1, axis=0)\n",
        "    core_ids_s = np.roll(core_ids, -1, axis=0)\n",
        "    core_ids_w = np.roll(core_ids, 1, axis=1)\n",
        "    core_ids_e = np.roll(core_ids, -1, axis=1)\n",
        "\n",
        "    self._permute_n = []\n",
        "    self._permute_s = []\n",
        "    self._permute_w = []\n",
        "    self._permute_e = []\n",
        "    for i, j in itertools.product(range(tpu_x), range(tpu_y)):\n",
        "      self._permute_n.append([core_ids_n[i, j], core_ids[i, j]])\n",
        "      self._permute_s.append([core_ids_s[i, j], core_ids[i, j]])\n",
        "      self._permute_w.append([core_ids_w[i, j], core_ids[i, j]])\n",
        "      self._permute_e.append([core_ids_e[i, j], core_ids[i, j]])\n",
        "\n",
        "  def get_boundary_n(self, super_grids, i, j, black):\n",
        "    \"\"\"The boundary on the northern direction.\"\"\"\n",
        "    # Collective permute can handle single core condition naturally.\n",
        "    if black:\n",
        "      if i == 0:\n",
        "        boundary_n = tpu_ops.collective_permute(\n",
        "            super_grids[1][0][-1][j][-1, -1, :, :], self._permute_n)\n",
        "      else:\n",
        "        boundary_n = super_grids[1][0][i - 1][j][-1, -1, :, :]\n",
        "    else:\n",
        "      if i == 0:\n",
        "        # In this case, its northern boundary is the southern boundary of the\n",
        "        # sub-lattice in the replica above, assuming periodic boundary\n",
        "        # condition on core topology.\n",
        "        boundary_n = tpu_ops.collective_permute(\n",
        "            super_grids[1][1][-1][j][-1, -1, :, :], self._permute_n)\n",
        "      else:\n",
        "        boundary_n = super_grids[1][1][i - 1][j][-1, -1, :, :]\n",
        "\n",
        "    return boundary_n\n",
        "\n",
        "  def get_boundary_s(self, super_grids, i, j, black):\n",
        "    \"\"\"The boundary on the southern direction.\"\"\"\n",
        "    if black:\n",
        "      if i == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "        boundary_s = tpu_ops.collective_permute(\n",
        "            super_grids[0][1][0][j][0, 0, :, :], self._permute_s)\n",
        "      else:\n",
        "        boundary_s = super_grids[0][1][i + 1][j][0, 0, :, :]\n",
        "    else:\n",
        "      if i == LATTICE_SUPER_BLOCK_DIM -1:\n",
        "        # In this case, its southern boundary is the northern boundary of the\n",
        "        # sub-lattice in the replica below, assuming periodic boundary\n",
        "        # condition on core topology.\n",
        "        boundary_s = tpu_ops.collective_permute(\n",
        "            super_grids[0][0][0][j][0, 0, :, :], self._permute_s)\n",
        "      else:\n",
        "        boundary_s = super_grids[0][0][i + 1][j][0, 0, :, :]\n",
        "    return boundary_s\n",
        "\n",
        "  def get_boundary_w(self, super_grids, i, j, black):\n",
        "    \"\"\"The boundary on the western direction.\"\"\"\n",
        "    if black:\n",
        "      if j == 0:\n",
        "        boundary_w = tpu_ops.collective_permute(\n",
        "            super_grids[0][1][i][-1][:, :, -1, -1], self._permute_w)\n",
        "      else:\n",
        "        boundary_w = super_grids[0][1][i][j - 1][:, :, -1, -1]\n",
        "    else:\n",
        "      if j == 0:\n",
        "        # In this case, its western boundary is the eastern boundary of the\n",
        "        # sub-lattice in the replica on the left, assuming periodic boundary\n",
        "        # condition on core topology.\n",
        "        boundary_w = tpu_ops.collective_permute(\n",
        "            super_grids[1][1][i][-1][:, :, -1, -1], self._permute_w)\n",
        "      else:\n",
        "        boundary_w = super_grids[1][1][i][j - 1][:, :, -1, -1]\n",
        "    return boundary_w\n",
        "\n",
        "  def get_boundary_e(self, super_grids, i, j, black):\n",
        "    \"\"\"The boundary on the eastern direction.\"\"\"\n",
        "    if black:\n",
        "      if j == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "        boundary_e = tpu_ops.collective_permute(\n",
        "            super_grids[1][0][i][0][:, :, 0, 0], self._permute_e)\n",
        "      else:\n",
        "        boundary_e = super_grids[1][0][i][j + 1][:, :, 0, 0]\n",
        "    else:\n",
        "      if j == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "        # In this case, its eastern boundary is the western boundary of the\n",
        "        # sub-lattice in the replica on the right, assuming periodic boundary\n",
        "        # condition on core topology.\n",
        "        boundary_e = tpu_ops.collective_permute(\n",
        "            super_grids[0][0][i][0][:, :, 0, 0], self._permute_e)\n",
        "      else:\n",
        "        boundary_e = super_grids[0][0][i][j + 1][:, :, 0, 0]\n",
        "    return boundary_e\n",
        "\n",
        "  def sum_of_nearest_neighbors_black(self, super_grids):\n",
        "    \"\"\"The sum of nearest neighbor in each site on periodic boundaries.\"\"\"\n",
        "    sum_nn_00 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    sum_nn_11 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      neighbor_n_00 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', tf.einsum(\n",
        "              'i,j-\u003eij', tf.one_hot(0, KERNEL_CHANNEL_SIZE_2,\n",
        "                                    dtype=get_dtype()),\n",
        "              tf.one_hot(0, CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype())),\n",
        "              self.get_boundary_n(super_grids, i, j, black=True))\n",
        "      neighbor_w_00 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', self.get_boundary_w(super_grids, i, j, black=True),\n",
        "          tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot(0, CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE,\n",
        "                         dtype=get_dtype()),\n",
        "              tf.one_hot(0, KERNEL_CHANNEL_SIZE, dtype=get_dtype())))\n",
        "      neighbor_s_11 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot(KERNEL_CHANNEL_SIZE_2 - 1, KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype()),\n",
        "              tf.one_hot((CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2) - 1,\n",
        "                         CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype())),\n",
        "                         self.get_boundary_s(super_grids, i, j, black=True))\n",
        "      neighbor_e_11 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', self.get_boundary_e(super_grids, i, j, black=True),\n",
        "          tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot((CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE) - 1,\n",
        "              CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE, dtype=get_dtype()),\n",
        "              tf.one_hot(KERNEL_CHANNEL_SIZE - 1, KERNEL_CHANNEL_SIZE,\n",
        "                         dtype=get_dtype())))\n",
        "\n",
        "      sum_nn_00[i][j] = (tf.nn.conv2d(\n",
        "          super_grids[0][1][i][j], self._kernel_nn, padding='SAME') +\n",
        "          tf.einsum('ijkl-\u003elkji', tf.nn.conv2d(\n",
        "              tf.einsum('ijkl-\u003elkji', super_grids[1][0][i][j]),\n",
        "              self._kernel_nn_2, padding='SAME')))\n",
        "      sum_nn_11[i][j] = (tf.nn.conv2d(\n",
        "          super_grids[1][0][i][j], self._kernel_nn_t, padding='SAME') +\n",
        "            tf.einsum('ijkl-\u003elkji', tf.nn.conv2d(\n",
        "                tf.einsum('ijkl-\u003elkji', super_grids[0][1][i][j]),\n",
        "                self._kernel_nn_t_2, padding='SAME')))\n",
        "\n",
        "      sum_nn_00[i][j] += (neighbor_n_00 + neighbor_w_00)\n",
        "      sum_nn_11[i][j] += (neighbor_s_11 + neighbor_e_11)\n",
        "    return sum_nn_00, sum_nn_11\n",
        "\n",
        "  def sum_of_nearest_neighbors_white(self, super_grids):\n",
        "    \"\"\"The sum of nearest neighbor in each site on periodic boundaries.\"\"\"\n",
        "    sum_nn_01 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    sum_nn_10 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      neighbor_n_01 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', tf.einsum(\n",
        "              'i,j-\u003eij', tf.one_hot(0, KERNEL_CHANNEL_SIZE_2,\n",
        "                                    dtype=get_dtype()),\n",
        "              tf.one_hot(0, CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype())),\n",
        "              self.get_boundary_n(super_grids, i, j, black=False))\n",
        "      neighbor_e_01 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', self.get_boundary_e(super_grids, i, j, black=False),\n",
        "          tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot((CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE) - 1,\n",
        "              CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE, dtype=get_dtype()),\n",
        "              tf.one_hot(KERNEL_CHANNEL_SIZE - 1, KERNEL_CHANNEL_SIZE,\n",
        "                         dtype=get_dtype())))\n",
        "      neighbor_s_10 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot(KERNEL_CHANNEL_SIZE_2 - 1, KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype()),\n",
        "              tf.one_hot((CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2) - 1,\n",
        "                         CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2,\n",
        "                         dtype=get_dtype())),\n",
        "                         self.get_boundary_s(super_grids, i, j, black=False))\n",
        "      neighbor_w_10 = tf.einsum(\n",
        "          'ij,kl-\u003eijkl', self.get_boundary_w(super_grids, i, j, black=False),\n",
        "          tf.einsum(\n",
        "              'i,j-\u003eij',\n",
        "              tf.one_hot(0, CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE,\n",
        "                         dtype=get_dtype()),\n",
        "              tf.one_hot(0, KERNEL_CHANNEL_SIZE, dtype=get_dtype())))\n",
        "\n",
        "      sum_nn_01[i][j] = (tf.nn.conv2d(\n",
        "          super_grids[0][0][i][j], self._kernel_nn_t, padding='SAME') +\n",
        "          tf.einsum('ijkl-\u003elkji', tf.nn.conv2d(\n",
        "              tf.einsum('ijkl-\u003elkji', super_grids[1][1][i][j]),\n",
        "              self._kernel_nn_2, padding='SAME')))\n",
        "\n",
        "      sum_nn_10[i][j] = (tf.nn.conv2d(\n",
        "          super_grids[1][1][i][j], self._kernel_nn, padding='SAME') +\n",
        "            tf.einsum('ijkl-\u003elkji',\n",
        "                tf.nn.conv2d(\n",
        "                    tf.einsum('ijkl-\u003elkji', super_grids[0][0][i][j]),\n",
        "                    self._kernel_nn_t_2, padding='SAME')))\n",
        "\n",
        "      sum_nn_01[i][j] += (neighbor_n_01 + neighbor_e_01)\n",
        "      sum_nn_10[i][j] += (neighbor_s_10 + neighbor_w_10)\n",
        "    return sum_nn_01, sum_nn_10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qU9blJij6sXn"
      },
      "outputs": [],
      "source": [
        "class NearestNeighborCalculatorOptim(object):\n",
        "  \"\"\"Calculate the sum of nearest neighbor spin values.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    # Constant matrix to compute the sum of nearest neighbor spins. Applying\n",
        "    # it as a left operand to adding the northern/southern neighbors, as a\n",
        "    # right operand to adding the western/eastern neighobrs.\n",
        "    grid_nn = tf.constant(\n",
        "        np.eye(LATTICE_SUB_DIM // 2, k=0, dtype=np.float16) +\n",
        "        np.eye(LATTICE_SUB_DIM // 2, k=1, dtype=np.float16),\n",
        "        dtype=get_dtype())\n",
        "    self._grid_nn = tf.broadcast_to(grid_nn, [\n",
        "        LATTICE_BLOCK_DIMS[0],\n",
        "        LATTICE_BLOCK_DIMS[1],\n",
        "        LATTICE_SUB_DIM // 2,\n",
        "        LATTICE_SUB_DIM // 2,\n",
        "    ])\n",
        "    # Given the spins on northern/southern or western/eastern boundaries of a\n",
        "    # grid of sub-lattices, transform them into another grid of sub-lattices\n",
        "    # that are added to compensate the nearest neighbor sums for spins on the\n",
        "    # boundaries.\n",
        "    self._grid_expand_s = tf.broadcast_to(\n",
        "        tf.one_hot(0, LATTICE_SUB_DIM // 2, dtype=get_dtype()), [\n",
        "            LATTICE_BLOCK_DIMS[0],\n",
        "            LATTICE_BLOCK_DIMS[1],\n",
        "            LATTICE_SUB_DIM // 2,\n",
        "        ])\n",
        "    self._grid_expand_e = tf.broadcast_to(\n",
        "        tf.one_hot(\n",
        "            LATTICE_SUB_DIM // 2 - 1, LATTICE_SUB_DIM // 2, dtype=get_dtype()),\n",
        "        [\n",
        "            LATTICE_BLOCK_DIMS[0],\n",
        "            LATTICE_BLOCK_DIMS[1],\n",
        "            LATTICE_SUB_DIM // 2,\n",
        "        ])\n",
        "\n",
        "    # If the lattice is distributed in multiple replicas, then define that\n",
        "    # permutation pairs that permute the boundaries of sub-lattices across\n",
        "    # replicas in all 4 directions. Those boundary spins are used to compute the\n",
        "    # nearest neighor sums of boundary spins on sub-lattices. They are ignored\n",
        "    # if the lattice is updated on single core.\n",
        "    tpu_x, tpu_y = CORES_TOPOLOGY\n",
        "    core_ids = np.arange(tpu_x * tpu_y).reshape((tpu_x, tpu_y))\n",
        "    core_ids_n = np.roll(core_ids, 1, axis=0)\n",
        "    core_ids_s = np.roll(core_ids, -1, axis=0)\n",
        "    core_ids_w = np.roll(core_ids, 1, axis=1)\n",
        "    core_ids_e = np.roll(core_ids, -1, axis=1)\n",
        "\n",
        "    self._permute_n = []\n",
        "    self._permute_s = []\n",
        "    self._permute_w = []\n",
        "    self._permute_e = []\n",
        "    for i, j in itertools.product(range(tpu_x), range(tpu_y)):\n",
        "      self._permute_n.append([core_ids_n[i, j], core_ids[i, j]])\n",
        "      self._permute_s.append([core_ids_s[i, j], core_ids[i, j]])\n",
        "      self._permute_w.append([core_ids_w[i, j], core_ids[i, j]])\n",
        "      self._permute_e.append([core_ids_e[i, j], core_ids[i, j]])\n",
        "\n",
        "  def get_boundary_n(self, super_grids, i, j, single_core, black):\n",
        "    \"\"\"The boundary on the northern direction.\"\"\"\n",
        "    if single_core:\n",
        "      if black:\n",
        "        boundary_n = super_grids[1][0][\n",
        "            (i - 1) % LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :]\n",
        "        boundary_n_rest = super_grids[1][0][i][j][:-1, :, -1, :]\n",
        "      else:\n",
        "        boundary_n = super_grids[1][1][\n",
        "            (i - 1) % LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :]\n",
        "        boundary_n_rest = super_grids[1][1][i][j][:-1, :, -1, :]\n",
        "    else:\n",
        "      if black:\n",
        "        if i == 0:\n",
        "          boundary_n = tpu_ops.collective_permute(\n",
        "              super_grids[1][0][(i - 1) %\n",
        "                                LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :],\n",
        "              self._permute_n)\n",
        "        else:\n",
        "          boundary_n = super_grids[1][0][\n",
        "              (i - 1) % LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :]\n",
        "        boundary_n_rest = super_grids[1][0][i][j][:-1, :, -1, :]\n",
        "      else:\n",
        "        if i == 0:\n",
        "          # In this case, its northern boundary is the southern boundary of the\n",
        "          # sub-lattice in the replica above, assuming periodic boundary\n",
        "          # condition on core topology.\n",
        "          boundary_n = tpu_ops.collective_permute(\n",
        "              super_grids[1][1][(i - 1) %\n",
        "                                LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :],\n",
        "              self._permute_n)\n",
        "        else:\n",
        "          boundary_n = super_grids[1][1][\n",
        "              (i - 1) % LATTICE_SUPER_BLOCK_DIM][j][-1:, :, -1, :]\n",
        "        boundary_n_rest = super_grids[1][1][i][j][:-1, :, -1, :]\n",
        "    grid_boundary_n_ij = tf.concat([boundary_n, boundary_n_rest], axis=0)\n",
        "    return grid_boundary_n_ij\n",
        "\n",
        "  def get_boundary_s(self, super_grids, i, j, single_core, black):\n",
        "    \"\"\"The boundary on the southern direction.\"\"\"\n",
        "    if single_core:\n",
        "      if black:\n",
        "        boundary_s = super_grids[0][1][(i + 1) %\n",
        "                                       LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :]\n",
        "        boundary_s_rest = super_grids[0][1][i][j][1:, :, 0, :]\n",
        "      else:\n",
        "        boundary_s = super_grids[0][0][(i + 1) %\n",
        "                                       LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :]\n",
        "        boundary_s_rest = super_grids[0][0][i][j][1:, :, 0, :]\n",
        "    else:\n",
        "      if black:\n",
        "        if i == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "          boundary_s = tpu_ops.collective_permute(\n",
        "              super_grids[0][1][(i + 1) %\n",
        "                                LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :],\n",
        "              self._permute_s)\n",
        "        else:\n",
        "          boundary_s = super_grids[0][1][\n",
        "              (i + 1) % LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :]\n",
        "        boundary_s_rest = super_grids[0][1][i][j][1:, :, 0, :]\n",
        "      else:\n",
        "        if i == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "          # In this case, its southern boundary is the northern boundary of the\n",
        "          # sub-lattice in the replica below, assuming periodic boundary\n",
        "          # condition on core topology.\n",
        "          boundary_s = tpu_ops.collective_permute(\n",
        "              super_grids[0][0][(i + 1) %\n",
        "                                LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :],\n",
        "              self._permute_s)\n",
        "        else:\n",
        "          boundary_s = super_grids[0][0][\n",
        "              (i + 1) % LATTICE_SUPER_BLOCK_DIM][j][:1, :, 0, :]\n",
        "        boundary_s_rest = super_grids[0][0][i][j][1:, :, 0, :]\n",
        "    grid_boundary_s_ij = tf.concat([boundary_s_rest, boundary_s], axis=0)\n",
        "    return grid_boundary_s_ij\n",
        "\n",
        "  def get_boundary_w(self, super_grids, i, j, single_core, black):\n",
        "    \"\"\"The boundary on the western direction.\"\"\"\n",
        "    if single_core:\n",
        "      if black:\n",
        "        boundary_w = super_grids[0][1][i][\n",
        "            (j - 1) % LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1]\n",
        "        boundary_w_rest = super_grids[0][1][i][j][:, :-1, :, -1]\n",
        "      else:\n",
        "        boundary_w = super_grids[1][1][i][\n",
        "            (j - 1) % LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1]\n",
        "        boundary_w_rest = super_grids[1][1][i][j][:, :-1, :, -1]\n",
        "    else:\n",
        "      if black:\n",
        "        if j == 0:\n",
        "          boundary_w = tpu_ops.collective_permute(\n",
        "              super_grids[0][1][i][(j - 1) %\n",
        "                                   LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1],\n",
        "              self._permute_w)\n",
        "        else:\n",
        "          boundary_w = super_grids[0][1][i][\n",
        "              (j - 1) % LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1]\n",
        "        boundary_w_rest = super_grids[0][1][i][j][:, :-1, :, -1]\n",
        "      else:\n",
        "        if j == 0:\n",
        "          # In this case, its western boundary is the eastern boundary of the\n",
        "          # sub-lattice in the replica on the left, assuming periodic boundary\n",
        "          # condition on core topology.\n",
        "          boundary_w = tpu_ops.collective_permute(\n",
        "              super_grids[1][1][i][(j - 1) %\n",
        "                                   LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1],\n",
        "              self._permute_w)\n",
        "        else:\n",
        "          boundary_w = super_grids[1][1][i][\n",
        "              (j - 1) % LATTICE_SUPER_BLOCK_DIM][:, -1:, :, -1]\n",
        "        boundary_w_rest = super_grids[1][1][i][j][:, :-1, :, -1]\n",
        "    grid_boundary_w_ij = tf.concat([boundary_w, boundary_w_rest], axis=1)\n",
        "    return grid_boundary_w_ij\n",
        "\n",
        "  def get_boundary_e(self, super_grids, i, j, single_core, black):\n",
        "    \"\"\"The boundary on the eastern direction.\"\"\"\n",
        "    if single_core:\n",
        "      if black:\n",
        "        boundary_e = super_grids[1][0][i][(j + 1) %\n",
        "                                          LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0]\n",
        "        boundary_e_rest = super_grids[1][0][i][j][:, 1:, :, 0]\n",
        "      else:\n",
        "        boundary_e = super_grids[0][0][i][(j + 1) %\n",
        "                                          LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0]\n",
        "        boundary_e_rest = super_grids[0][0][i][j][:, 1:, :, 0]\n",
        "    else:\n",
        "      if black:\n",
        "        if j == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "          boundary_e = tpu_ops.collective_permute(\n",
        "              super_grids[1][0][i][(j + 1) %\n",
        "                                   LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0],\n",
        "              self._permute_e)\n",
        "        else:\n",
        "          boundary_e = super_grids[1][0][i][\n",
        "              (j + 1) % LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0]\n",
        "        boundary_e_rest = super_grids[1][0][i][j][:, 1:, :, 0]\n",
        "      else:\n",
        "        if j == LATTICE_SUPER_BLOCK_DIM - 1:\n",
        "          # In this case, its eastern boundary is the western boundary of the\n",
        "          # sub-lattice in the replica on the right, assuming periodic boundary\n",
        "          # condition on core topology.\n",
        "          boundary_e = tpu_ops.collective_permute(\n",
        "              super_grids[0][0][i][(j + 1) %\n",
        "                                   LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0],\n",
        "              self._permute_e)\n",
        "        else:\n",
        "          boundary_e = super_grids[0][0][i][\n",
        "              (j + 1) % LATTICE_SUPER_BLOCK_DIM][:, :1, :, 0]\n",
        "        boundary_e_rest = super_grids[0][0][i][j][:, 1:, :, 0]\n",
        "    grid_boundary_e_ij = tf.concat([boundary_e_rest, boundary_e], axis=1)\n",
        "    return grid_boundary_e_ij\n",
        "\n",
        "  def sum_of_nearest_neighbors_black(self, super_grids, single_core):\n",
        "    \"\"\"The sum of nearest neighbor in each site on periodic boundaries.\"\"\"\n",
        "    sum_nn_00 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      sum_nn_00[i][j] = (\n",
        "          tf.matmul(super_grids[0][1][i][j], self._grid_nn) +\n",
        "          tf.matmul(self._grid_nn, super_grids[1][0][i][j], transpose_a=True))\n",
        "    sum_nn_11 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      sum_nn_11[i][j] = (\n",
        "          tf.matmul(self._grid_nn, super_grids[0][1][i][j]) +\n",
        "          tf.matmul(super_grids[1][0][i][j], self._grid_nn, transpose_b=True))\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      # Handle the northern/western boundary.\n",
        "      grid_boundary_n_ij = self.get_boundary_n(super_grids, i, j, single_core,\n",
        "                                               True)\n",
        "      grid_boundary_w_ij = self.get_boundary_w(super_grids, i, j, single_core,\n",
        "                                               True)\n",
        "\n",
        "      sum_nn_00[i][j] += (\n",
        "          tf.einsum('mni,mnj-\u003emnij', self._grid_expand_s, grid_boundary_n_ij) +\n",
        "          tf.einsum('mni,mnj-\u003emnij', grid_boundary_w_ij, self._grid_expand_s))\n",
        "\n",
        "      # Handle the southern/eastern boundary.\n",
        "      grid_boundary_s_ij = self.get_boundary_s(super_grids, i, j, single_core,\n",
        "                                               True)\n",
        "      grid_boundary_e_ij = self.get_boundary_e(super_grids, i, j, single_core,\n",
        "                                               True)\n",
        "      sum_nn_11[i][j] += (\n",
        "          tf.einsum('mni,mnj-\u003emnij', self._grid_expand_e, grid_boundary_s_ij) +\n",
        "          tf.einsum('mni,mnj-\u003emnij', grid_boundary_e_ij, self._grid_expand_e))\n",
        "\n",
        "    return sum_nn_00, sum_nn_11\n",
        "\n",
        "  def sum_of_nearest_neighbors_white(self, super_grids, single_core):\n",
        "    \"\"\"The sum of nearest neighbor in each site on periodic boundaries.\"\"\"\n",
        "    sum_nn_01 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      sum_nn_01[i][j] = (\n",
        "          tf.matmul(super_grids[0][0][i][j], self._grid_nn, transpose_b=True) +\n",
        "          tf.matmul(self._grid_nn, super_grids[1][1][i][j], transpose_a=True))\n",
        "    sum_nn_10 = create_list([LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      sum_nn_10[i][j] = (\n",
        "          tf.matmul(self._grid_nn, super_grids[0][0][i][j]) +\n",
        "          tf.matmul(super_grids[1][1][i][j], self._grid_nn))\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      # Handle the northern/estern boundary.\n",
        "      grid_boundary_n_ij = self.get_boundary_n(super_grids, i, j, single_core,\n",
        "                                               False)\n",
        "      grid_boundary_e_ij = self.get_boundary_e(super_grids, i, j, single_core,\n",
        "                                               False)\n",
        "      sum_nn_01[i][j] += (\n",
        "          tf.einsum('mni,mnj-\u003emnij', self._grid_expand_s, grid_boundary_n_ij) +\n",
        "          tf.einsum('mni,mnj-\u003emnij', grid_boundary_e_ij, self._grid_expand_e))\n",
        "      # Handle the southern/western boundary.\n",
        "      grid_boundary_s_ij = self.get_boundary_s(super_grids, i, j, single_core,\n",
        "                                               False)\n",
        "      grid_boundary_w_ij = self.get_boundary_w(super_grids, i, j, single_core,\n",
        "                                               False)\n",
        "      sum_nn_10[i][j] += (\n",
        "          tf.einsum('mni,mnj-\u003emnij', self._grid_expand_e, grid_boundary_s_ij) +\n",
        "          tf.einsum('mni,mnj-\u003emnij', grid_boundary_w_ij, self._grid_expand_s))\n",
        "\n",
        "    return sum_nn_01, sum_nn_10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "INFpR87M68F2"
      },
      "outputs": [],
      "source": [
        "def _validate_params():\n",
        "  \"\"\"Validate parameters before using them.\"\"\"\n",
        "  assert LATTICE_SUB_DIM \u003e 0\n",
        "\n",
        "  assert LATTICE_BLOCK_DIMS is not None and len(LATTICE_BLOCK_DIMS) == 2\n",
        "  assert LATTICE_BLOCK_DIMS[0] \u003e 0 and LATTICE_BLOCK_DIMS[1] \u003e 0\n",
        "  assert LATTICE_SUPER_BLOCK_DIM \u003e 0\n",
        "  assert NORMALIZED_TEMPERATURE \u003e 0.0\n",
        "  assert NUMBER_OF_BURN_IN_UPDATE \u003e 0\n",
        "  assert NUMBER_OF_WHOLE_LATTICE_UPDATE \u003e 0\n",
        "\n",
        "  assert CORES_TOPOLOGY is not None and len(CORES_TOPOLOGY) == 2\n",
        "  assert CORES_TOPOLOGY[0] \u003e 0 and CORES_TOPOLOGY[1] \u003e 0\n",
        "\n",
        "\n",
        "def compute_nanoseconds_per_flip(step_time,\n",
        "                                 num_steps=NUMBER_OF_WHOLE_LATTICE_UPDATE):\n",
        "  \"\"\"Compute the nanoseconds per flip given step_time in seconds.\"\"\"\n",
        "  nanos_per_flip = step_time * _NANOS_PER_SECOND / (\n",
        "      (LATTICE_SUB_DIM * LATTICE_SUPER_BLOCK_DIM)**2 *\n",
        "      np.prod(LATTICE_BLOCK_DIMS) * num_steps *\n",
        "      np.prod(CORES_TOPOLOGY))\n",
        "  return nanos_per_flip\n",
        "\n",
        "\n",
        "# pylint: disable=unused-argument\n",
        "def _grid_initializer(shape, dtype, partition_info):\n",
        "  grid_value_init = 2.0 * tf.cast(\n",
        "      tf.random_uniform(shape) \u003e 0.5, dtype=dtype) - 1.0\n",
        "  return grid_value_init"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Y3mMuc7xaCar"
      },
      "outputs": [],
      "source": [
        "def update_optim_conv(sweeps, fn):\n",
        "  \"\"\"Simulation in each replica using 'compact' representation.\n",
        "\n",
        "  Boundary value communications are handled by collective permute. This uses\n",
        "  the convolution instead of matmul for the nearest neighbor sum calculation.\n",
        "\n",
        "  Ref: Algorithm UpdateOptim in the paper.\n",
        "\n",
        "  Args:\n",
        "    sweeps: the number of whole lattice update.\n",
        "    fn: the function on a given configuration of the lattice, it takes a 2-D\n",
        "      list of variables, i.e., v_ij and returns a float32.\n",
        "\n",
        "  Returns:\n",
        "    The estimated expectation of fn, i.e., \u003cfn\u003e\n",
        "  \"\"\"\n",
        "  _validate_params()\n",
        "\n",
        "  temperature_muliplier = -2 * get_inverse_temperature()\n",
        "\n",
        "  lattice_sub_shape = [\n",
        "      KERNEL_CHANNEL_SIZE_2,\n",
        "      CONV_LATTICE_DIMS[0] // KERNEL_CHANNEL_SIZE_2,\n",
        "      CONV_LATTICE_DIMS[1] // KERNEL_CHANNEL_SIZE,\n",
        "      KERNEL_CHANNEL_SIZE\n",
        "  ]\n",
        "\n",
        "  nn_calculator = NearestNeighborCalculatorOptimConv()\n",
        "\n",
        "  def checkerboard(k, estimation):\n",
        "    \"\"\"Checkerboard algorithm update.\"\"\"\n",
        "    with tf.variable_scope('tpu', reuse=tf.AUTO_REUSE, use_resource=True):\n",
        "      super_grids = create_list([2] * 2 + [LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "      for l, m in create_iterator([2] * 2):\n",
        "        for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "          super_grids[l][m][i][j] = tf.get_variable(\n",
        "           'grids_%d%d_%d%d' % (l, m, i, j),\n",
        "           initializer=_grid_initializer,\n",
        "           shape=lattice_sub_shape,\n",
        "           dtype=tf.bfloat16)\n",
        "\n",
        "    def update(probs, black):\n",
        "      \"\"\"Checkerboard algorithm update for a given color.\"\"\"\n",
        "      if black:\n",
        "        idx = [[0, 0], [1, 1]]\n",
        "        sum_nn_color = nn_calculator.sum_of_nearest_neighbors_black(super_grids)\n",
        "      else:\n",
        "        idx = [[0, 1], [1, 0]]\n",
        "        sum_nn_color = nn_calculator.sum_of_nearest_neighbors_white(super_grids)\n",
        "      assign_ops = []\n",
        "      for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "        for [idx0, idx1], sum_nn in zip(idx, sum_nn_color):\n",
        "          acceptance_ratio_ij = (\n",
        "            temperature_muliplier * sum_nn[i][j] *\n",
        "            super_grids[idx0][idx1][i][j])\n",
        "          flips_ij = tf.cast(\n",
        "            probs[idx0][idx1][i][j] \u003c acceptance_ratio_ij, dtype=get_dtype())\n",
        "          assign_ops.append(super_grids[idx0][idx1][i][j].assign_sub(\n",
        "                flips_ij * super_grids[idx0][idx1][i][j] *\n",
        "                tf.constant(2.0, dtype=get_dtype())))\n",
        "      return assign_ops\n",
        "\n",
        "    probs = create_list([2] * 2 + [LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for l, m in create_iterator([2] * 2):\n",
        "      for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "        probs[l][m][i][j] = (\n",
        "          tf.log(tf.random_uniform(lattice_sub_shape, dtype=get_dtype())))\n",
        "    grid_black_update = update(probs, black=True)\n",
        "    with tf.control_dependencies(grid_black_update):\n",
        "      grid_white_update = update(probs, black=False)\n",
        "    with tf.control_dependencies(grid_white_update):\n",
        "      return k + 1, (\n",
        "          estimation * tf.cast(k, dtype=tf.float32) /\n",
        "          tf.cast(k + 1, dtype=tf.float32) +\n",
        "          tf.cast(fn(super_grids), dtype=tf.float32) /\n",
        "          tf.cast(k + 1, dtype=tf.float32))\n",
        "\n",
        "  def while_loop(sweeps):\n",
        "    _, estimation = tf.while_loop(lambda i, _: i \u003c sweeps, checkerboard, [\n",
        "        0,\n",
        "        tf.constant([0.0, 0.0, 0.0], dtype=tf.float32),\n",
        "    ])\n",
        "    return estimation\n",
        "\n",
        "  return while_loop(sweeps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "truzk50Z7BQA"
      },
      "outputs": [],
      "source": [
        "def update_optim(sweeps, fn, single_core):\n",
        "  \"\"\"Simulation in each replica using 'compact' representation.\n",
        "\n",
        "  Boundary value communications are handled by collective permute.\n",
        "\n",
        "  Ref: Algorithm UpdateOptim in the paper.\n",
        "\n",
        "  Args:\n",
        "    sweeps: the number of whole lattice update.\n",
        "    fn: the function on a given configuration of the lattice, it takes a 2-D\n",
        "      list of variables, i.e., v_ij and returns a float32.\n",
        "    single_core: a bool that specify whether the lattice is updated on single or\n",
        "      multiple cores.\n",
        "\n",
        "  Returns:\n",
        "    The estimated expectation of fn, i.e., \u003cfn\u003e\n",
        "  \"\"\"\n",
        "  _validate_params()\n",
        "\n",
        "  temperature_muliplier = -2 * get_inverse_temperature()\n",
        "\n",
        "  lattice_sub_shape = [\n",
        "      LATTICE_BLOCK_DIMS[0],\n",
        "      LATTICE_BLOCK_DIMS[1],\n",
        "      LATTICE_SUB_DIM // 2,\n",
        "      LATTICE_SUB_DIM // 2,\n",
        "  ]\n",
        "\n",
        "  nn_calculator = NearestNeighborCalculatorOptim()\n",
        "\n",
        "  def checkerboard(k, estimation):\n",
        "    \"\"\"Checkerboard algorithm update.\"\"\"\n",
        "    with tf.variable_scope('tpu', reuse=tf.AUTO_REUSE, use_resource=True):\n",
        "      super_grids = create_list([2] * 2 + [LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "      for l, m in create_iterator([2] * 2):\n",
        "        for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "          super_grids[l][m][i][j] = tf.get_variable(\n",
        "              'grids_%d%d_%d%d' % (l, m, i, j),\n",
        "              initializer=_grid_initializer,\n",
        "              shape=lattice_sub_shape,\n",
        "              dtype=tf.bfloat16)\n",
        "\n",
        "    def update(probs, black):\n",
        "      \"\"\"Checkerboard algorithm update for a given color.\"\"\"\n",
        "      if black:\n",
        "        idx = [[0, 0], [1, 1]]\n",
        "        sum_nn_color = \\\n",
        "            nn_calculator.sum_of_nearest_neighbors_black(\n",
        "                super_grids, single_core)\n",
        "      else:\n",
        "        idx = [[0, 1], [1, 0]]\n",
        "        sum_nn_color = \\\n",
        "            nn_calculator.sum_of_nearest_neighbors_white(\n",
        "                super_grids, single_core)\n",
        "      assign_ops = []\n",
        "      for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "        for [idx0, idx1], sum_nn in zip(idx, sum_nn_color):\n",
        "          acceptance_ratio_ij = (\n",
        "              temperature_muliplier * sum_nn[i][j] *\n",
        "              super_grids[idx0][idx1][i][j])\n",
        "          flips_ij = tf.cast(\n",
        "              probs[idx0][idx1][i][j] \u003c acceptance_ratio_ij, dtype=get_dtype())\n",
        "          assign_ops.append(super_grids[idx0][idx1][i][j].assign_sub(\n",
        "              flips_ij * super_grids[idx0][idx1][i][j] *\n",
        "              tf.constant(2.0, dtype=get_dtype())))\n",
        "      return assign_ops\n",
        "\n",
        "    probs = create_list([2] * 2 + [LATTICE_SUPER_BLOCK_DIM] * 2)\n",
        "    for l, m in create_iterator([2] * 2):\n",
        "      for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "        probs[l][m][i][j] = (\n",
        "            tf.log(tf.random_uniform(lattice_sub_shape, dtype=get_dtype())))\n",
        "    grid_black_update = update(probs, black=True)\n",
        "    with tf.control_dependencies(grid_black_update):\n",
        "      grid_white_update = update(probs, black=False)\n",
        "    with tf.control_dependencies(grid_white_update):\n",
        "      return k + 1, (\n",
        "          estimation * tf.cast(k, dtype=tf.float32) /\n",
        "          tf.cast(k + 1, dtype=tf.float32) +\n",
        "          tf.cast(fn(super_grids), dtype=tf.float32) /\n",
        "          tf.cast(k + 1, dtype=tf.float32))\n",
        "\n",
        "  def while_loop(sweeps):\n",
        "    _, estimation = tf.while_loop(lambda i, _: i \u003c sweeps, checkerboard, [\n",
        "        0,\n",
        "        tf.constant([0.0, 0.0, 0.0], dtype=tf.float32),\n",
        "    ])\n",
        "    return estimation\n",
        "\n",
        "  return while_loop(sweeps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a2XFjMu0Yjtk"
      },
      "outputs": [],
      "source": [
        "def create_ising_mcmc_simulator(fn, single_core, device_assignment):\n",
        "  \"\"\"Ising model MCMC simulation on single or multiple TPUs.\n",
        "\n",
        "  The whole lattice is distributed equally among TPU cores, and each TPU core\n",
        "  runs one replica given cores topology `[l1, l2]`. In each core, because of\n",
        "  the hard limit of protobuf (2GB), we split the sub-lattice into multi-scale\n",
        "  sub-lattices as follows:\n",
        "\n",
        "  The sub-lattice in each core is divided into a 'kxk' grid of sub-lattices,\n",
        "  each sub-lattice is a tensor variable, where `k' is flag:\n",
        "  lattice_super_block_dim,\n",
        "  `[v_00, ........... v_{0,k-1}]`\n",
        "  `[v_10, ........... v_{1,k-1}]`\n",
        "  `.............................`\n",
        "  `[v_{k-1,0}, ..., v_{k-1,k-1}]`\n",
        "\n",
        "  Each v_ij is again a `m1xm2` grid of smaller sub-lattice, where `m1` and\n",
        "  `m2` are provided by flag: lattice_block_dims,\n",
        "  `[g_00, ............, g_{0,m2-1}]`\n",
        "  `[g_10, ............, g_{1,m2-1}]`\n",
        "  `................................`\n",
        "  `[g_{m1-1,0}, ..., g_{m1-1,m2-1}]`\n",
        "\n",
        "  and each g_ij is a `nxn` sub-lattice, where `n` is flag: lattice_sub_dim.\n",
        "\n",
        "  Each g_ij is furthure split into 4 compact sub-lattices, i.e.,\n",
        "  `g_ij_00 = g_ij[0::2, 0::2]`\n",
        "  `g_ij_01 = g_ij[0::2, 1::2]`\n",
        "  `g_ij_10 = g_ij[1::2, 0::2]`\n",
        "  `g_ij_11 = g_ij[1::2, 1::2]`\n",
        "\n",
        "  Where g_ij_00 and g_ij_11 are all black spins, and g_ij_01 and g_ij_10 are\n",
        "  all white spins.\n",
        "\n",
        "  Thus the whole lattice has dimensions [l1*k*m1*n, l2*k*m2*n].\n",
        "\n",
        "  The boundaries of each sub-lattice are collectively permuted among replicas\n",
        "  to calculate nearest neighbor sums, which are used to compute\n",
        "  Metropolis-Hastings updates.\n",
        "\n",
        "  Args:\n",
        "    fn: the function on a given configuration of the sub_lattice on the\n",
        "      replica. It takes a 2-D list of variables, i.e., v_ij, and returns a\n",
        "      float32. fn must be additive in order to insure the correctness of the\n",
        "      estimation, i.e., `fn([sub_lattice_1, ..., sub_lattice_n]) = sum_{i=1}^n\n",
        "      fn([sub_lattice_i])`.\n",
        "    single_core: a bool that specify whether the lattice is updated on single\n",
        "      or multiple cores.\n",
        "    device_assignment: a `tf.contrib.tpu.DeviceAssignment` specifing the mapping\n",
        "      from replica id to the TPU mesh grid coordinate.\n",
        "\n",
        "  Returns:\n",
        "    The tuple of a `int32` placeholder for the number of flips and a handle\n",
        "    for the simulation computation.\n",
        "  \"\"\"\n",
        "  num_flips_placeholder = tf.placeholder(tf.int32, [])\n",
        "  def compute_fn(sweeps):\n",
        "    return (update_optim_conv(sweeps, fn) if USE_CONV else\n",
        "            update_optim(sweeps, fn, single_core))\n",
        "  return num_flips_placeholder, tf.contrib.tpu.replicate(\n",
        "      compute_fn,\n",
        "      inputs=[[num_flips_placeholder] for _ in range(np.prod(CORES_TOPOLOGY))],\n",
        "      device_assignment=device_assignment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PtIb7bvf7aeN"
      },
      "outputs": [],
      "source": [
        "def ising_mcmc_simulation(fn, session):\n",
        "  \"\"\"Ising mcmc simulation.\n",
        "\n",
        "  Args:\n",
        "    fn: the function on a given configuration of the lattice. It takes a 2-D\n",
        "      list of variables, i.e., v_ij and returns a float32. Refer to\n",
        "      create_ising_mcmc_on_single_core for details on the variables.\n",
        "    session: a function returns a tensorflow session.\n",
        "\n",
        "  Returns:\n",
        "    The estimation of the expectation of fn.\n",
        "  \"\"\"\n",
        "  with session() as sess:\n",
        "    topology = tf.contrib.tpu.Topology(\n",
        "        sess.run(tf.contrib.tpu.initialize_system()))\n",
        "    device_assignment, _ = tpu_device_assignment(\n",
        "        np.asarray(CORES_TOPOLOGY), topology)\n",
        "    single_core = is_single_core(CORES_TOPOLOGY)\n",
        "    print(sess.list_devices())\n",
        "    num_flips_placeholder, compute_handle = create_ising_mcmc_simulator(\n",
        "        fn, single_core, device_assignment)\n",
        "    sess.run(tf.global_variables_initializer())\n",
        "    feed_dict = {num_flips_placeholder: NUMBER_OF_BURN_IN_UPDATE}\n",
        "    print('--- start burning in ---')\n",
        "    burn_in_start_time = time.time()\n",
        "    sess.run(compute_handle, feed_dict=feed_dict)\n",
        "    burn_in_time = time.time() - burn_in_start_time\n",
        "    print('--- finish burning in ---')\n",
        "    print('--- burn-in time: %s seconds ---' % burn_in_time)\n",
        "    feed_dict = {num_flips_placeholder: NUMBER_OF_WHOLE_LATTICE_UPDATE}\n",
        "    start_time = time.time()\n",
        "    estimation_val = sess.run(compute_handle, feed_dict=feed_dict)\n",
        "    step_time = time.time() - start_time\n",
        "    nanos_per_flip = compute_nanoseconds_per_flip(step_time)\n",
        "    print('--- run time with overhead: %s seconds ---' % step_time)\n",
        "    print('--- step time with overhead: %s seconds per whole-lattice '\n",
        "          'flip ---' % (step_time /  NUMBER_OF_WHOLE_LATTICE_UPDATE))\n",
        "    print('--- %s nanoseconds per flip (with overhead)---' % nanos_per_flip)\n",
        "    step_time_no_overhead = step_time - burn_in_time\n",
        "    nanos_per_flip_no_overhead = compute_nanoseconds_per_flip(\n",
        "        step_time_no_overhead,\n",
        "        NUMBER_OF_WHOLE_LATTICE_UPDATE - NUMBER_OF_BURN_IN_UPDATE)\n",
        "    print('--- run time excluding overhead: %s seconds ---' %\n",
        "          step_time_no_overhead)\n",
        "    print('--- step time excluding overhead: %s seconds '\n",
        "          'per whole-lattice flip ---' % (step_time_no_overhead /\n",
        "          (NUMBER_OF_WHOLE_LATTICE_UPDATE - NUMBER_OF_BURN_IN_UPDATE)))\n",
        "    print('--- %s nanoseconds per flip (without overhead)---' %\n",
        "          nanos_per_flip_no_overhead)\n",
        "  return estimation_val"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "rQlj9Lxu997v"
      },
      "outputs": [],
      "source": [
        "def get_session():\n",
        "  def _get_tpu_setup():\n",
        "    tpu_cluster_resolver = TPUClusterResolver(tpu=os.environ['TPU_NAME'])\n",
        "    cluster_def = tpu_cluster_resolver.cluster_spec().as_cluster_def()\n",
        "    tpu_master_grpc_path = tpu_cluster_resolver.get_master()\n",
        "    return cluster_def, tpu_master_grpc_path\n",
        "\n",
        "  cluster_def, tpu_master_grpc_path = _get_tpu_setup()\n",
        "  config = tf.ConfigProto(\n",
        "      allow_soft_placement=True,\n",
        "      isolate_session_state=True,\n",
        "      cluster_def=cluster_def)\n",
        "  return tf.Session(tpu_master_grpc_path, config=config)\n",
        "\n",
        "\n",
        "def reduce_mean(super_grids):\n",
        "  mag = tf.constant(0.0, dtype=tf.float32)\n",
        "  num_base = np.prod([\n",
        "      2, 2, LATTICE_SUPER_BLOCK_DIM, LATTICE_SUPER_BLOCK_DIM\n",
        "  ]).astype(np.float32)\n",
        "  for l, m in create_iterator([2] * 2):\n",
        "    for i, j in create_iterator([LATTICE_SUPER_BLOCK_DIM] * 2):\n",
        "      mag += tf.cast(\n",
        "          tf.reduce_mean(super_grids[l][m][i][j]), dtype=tf.float32)\n",
        "  mag = mag / num_base\n",
        "  mag_sq = mag * mag\n",
        "  mag_4 = mag_sq * mag_sq\n",
        "\n",
        "  return tf.stack([mag, mag_sq, mag_4])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gD0UDPBgPMYH"
      },
      "source": [
        "Compute magnetization and Binder parameter for a given temperature"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XnD_IWEdOyeI"
      },
      "outputs": [],
      "source": [
        "NORMALIZED_TEMPERATURE = 0.1\n",
        "with tf.Graph().as_default():\n",
        "  estimation_val = ising_mcmc_simulation(reduce_mean, get_session)\n",
        "  print('estimations: %s' % estimation_val[0][0])\n",
        "  print('magnetization: %s' % estimation_val[0][0][0])\n",
        "  print('Binder parameter: %s' % (\n",
        "        1 - estimation_val[0][0][2] / (3 * estimation_val[0][0][1] ** 2)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IQ_dvio_Mpoz"
      },
      "source": [
        "Compute magnetization and Binder parameters for a range of temperatures"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wPuqEpVWfRcU"
      },
      "outputs": [],
      "source": [
        "results = []\n",
        "for t in np.arange(0.9, 1.1, 0.1):\n",
        "  NORMALIZED_TEMPERATURE = t\n",
        "  with tf.Graph().as_default():\n",
        "    estimation_val = ising_mcmc_simulation(reduce_mean, get_session)\n",
        "    print('estimations: %s' % estimation_val[0][0])\n",
        "    print('magnetization: %s' % estimation_val[0][0][0])\n",
        "    binder_param = (1 -\n",
        "      estimation_val[0][0][2] / (3 * estimation_val[0][0][1] ** 2))\n",
        "    print('Binder parameter: %s' % binder_param)\n",
        "    results.append((t, estimation_val[0][0][0], binder_param))\n",
        "for r in results:\n",
        "  print('temperature: %s, magnetization: %s, Binder parameter: %s' %\n",
        "        (r[0], r[1], r[2]))\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/brain/python/client:colab_notebook",
        "kind": "shared"
      },
      "name": "Ising_Model_MCMC_Simulation_on_TPU_with_convol.ipynb",
      "provenance": [
        {
          "file_id": "12ZZY_dDJ12F52gt6CSicCD8TOnKtq15x",
          "timestamp": 1573252231497
        },
        {
          "file_id": "1Zo0RbVPtEV2GVRzyOAUhuziIzQ4TS8fk",
          "timestamp": 1566861625912
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
